# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.

# SPDX-License-Identifier: Apache-2.0

import time

import pytest
import torch
from loguru import logger

import ttnn
from models.demos.yolov4.common import load_torch_model
from models.demos.yolov4.tt.downsample5 import Down5
from models.demos.yolov4.tt.model_preprocessing import create_ds5_model_parameters
from tests.ttnn.utils_for_testing import assert_with_pcc


@pytest.mark.parametrize(
    "resolution",
    [
        (320, 320),
        (640, 640),
    ],
)
@pytest.mark.parametrize("device_params", [{"l1_small_size": 16384}], indirect=True)
def test_down5(device, reset_seeds, model_location_generator, resolution):
    torch.manual_seed(0)

    torch_model = load_torch_model(model_location_generator, module="down5")
    if resolution == (320, 320):
        torch_input = torch.randn((1, 512, 20, 20), dtype=torch.float)
    elif resolution == (640, 640):
        torch_input = torch.randn((1, 512, 40, 40), dtype=torch.float)
    else:
        raise ValueError(f"Unsupported resolution: {resolution}")
    ref = torch_model(torch_input)

    parameters = create_ds5_model_parameters(torch_model, torch_input, resolution, device)

    torch_input = torch_input.permute(0, 2, 3, 1)
    ttnn_input = ttnn.from_torch(torch_input, dtype=ttnn.bfloat16)
    ttnn_model = Down5(device, parameters, parameters.conv_args)
    result_ttnn = ttnn_model(ttnn_input)

    start_time = time.time()
    for x in range(2):
        result_ttnn = ttnn_model(ttnn_input)
    logger.info(f"Time taken: {time.time() - start_time}")

    result = ttnn.to_torch(result_ttnn)
    result = result.permute(0, 3, 1, 2)
    result = result.reshape(ref.shape)
    assert_with_pcc(result, ref, 0.89)  # PCC 0.89 - The PCC will improve once #3612 is resolved.
